package com.ww.exercise.coroutine.mid.q17

import kotlinx.coroutines.*
import kotlinx.coroutines.sync.Semaphore
import kotlinx.coroutines.sync.withPermit
import java.sql.Connection
import java.sql.DriverManager
import java.sql.PreparedStatement
import java.util.concurrent.atomic.AtomicInteger
import kotlin.random.Random
import kotlin.system.measureTimeMillis

// 模拟数据库连接池
class DatabaseConnectionPool(
    private val url: String, private val user: String, private val password: String, private val maxConnections: Int
) {
    private val connectionQueue = ArrayDeque<Connection>()
    private val semaphore = Semaphore(maxConnections)

    init {
        // 初始化连接池
        repeat(minOf(5, maxConnections)) {
            connectionQueue.add(createNewConnection())
        }
    }

    private fun createNewConnection(): Connection {
        return DriverManager.getConnection(url, user, password).apply {
            autoCommit = false // 禁用自动提交
        }
    }

    suspend fun <T> getConnection(block: (Connection) -> T): T {
        semaphore.withPermit {
            val connection = if (connectionQueue.isNotEmpty()) {
                connectionQueue.removeFirst()
            } else {
                createNewConnection()
            }

            return try {
                block(connection)
            } finally {
                if (!connection.isClosed) {
                    connectionQueue.addLast(connection)
                }
            }
        }
    }

    fun close() {
        connectionQueue.forEach {
            try {
                it.close()
            } catch (e: Exception) {
            }
        }
        connectionQueue.clear()
    }
}

// 用户数据类
data class User(
    val id: Int, val username: String, val email: String, val age: Int, val registrationDate: Long
)

// 进度记录器（用于断点续插）
class ProgressRecorder(private val processFile: String = "insert_process.txt") {
    // 原子变量确保多线程环境下的进度记录安全
    private val progress = AtomicInteger(0)

    init {
        // 尝试从文件读取上次进度
        try {
            java.io.File(processFile).readText().toIntOrNull()?.let {
                progress.set(it)
            }
        } catch (e: Exception) {
            // 文件不存在或读取失败，从0开始
            println("未找到进度记录或记录损坏，将从0开始插入")
        }
    }

    // 获取当前进度
    fun getCurrentProgress() = progress.get()

    // 更新进度（原子操作）
    fun updateProgress(newProgress: Int) {
        progress.set(newProgress)
        // 异步写入文件，避免阻塞插入操作
        GlobalScope.launch(Dispatchers.IO) {
            try {
                java.io.File(processFile).writeText(newProgress.toString())
            } catch (e: Exception) {
                println("进度记录写入失败：${e.message}")
            }
        }
    }

    // 重置进度
    fun reset() {
        progress.set(0)
        try {
            java.io.File(processFile).delete()
        } catch (e: Exception) {
        }
    }
}

// 数据生成器
object UserDataGenerator {
    private val random = Random

    fun generateUsers(startId: Int, count: Int): List<User> {
        return (startId until startId + count).map { id ->
            val username = "user_${id}_${randomString(8)}"
            val email = "${username}@163.com"
            val age = random.nextInt(18, 65)
            val registrationDate = System.currentTimeMillis() - random.nextLong(365 * 24 * 60 * 60 * 1000L)

            User(id, username, email, age, registrationDate)
        }
    }

    private fun randomString(length: Int): String {
        val chars = "abcdefghijklmnopqrstuvwxyz0123456789"
        return (1..length).map { chars[random.nextInt(chars.length)] }.joinToString("")
    }
}

// 数据库操作工具类
class UserDatabaseHelper(private val connectionPool: DatabaseConnectionPool) {
    // 单条插入
    suspend fun insertUser(user: User) {
        connectionPool.getConnection { connection ->
            val sql = """
                INSERT INTO users (id, username, email, age, registration_date)
                VALUES (?, ?, ?, ?, ?)
            """.trimIndent()

            connection.prepareStatement(sql).use { stmt ->
                setUserParameters(stmt, user)
                stmt.executeUpdate()
                connection.commit()
            }
        }
    }

    // 批量插入
    suspend fun batchInsertUsers(users: List<User>): Int {
        if (users.isEmpty()) return 0

        return connectionPool.getConnection { connection ->
            val sql = """
                INSERT INTO users (id, username, email, age, registration_date)
                VALUES (?, ?, ?, ?, ?)
            """.trimIndent()

            connection.prepareStatement(sql).use { stmt ->
                users.forEach { user ->
                    setUserParameters(stmt, user)
                    stmt.addBatch()
                }

                val results = stmt.executeBatch()
                connection.commit()
                results.sum()
            }
        }
    }

    private fun setUserParameters(stmt: PreparedStatement, user: User) {
        stmt.setInt(1, user.id)
        stmt.setString(2, user.username)
        stmt.setString(3, user.email)
        stmt.setInt(4, user.age)
        stmt.setLong(5, user.registrationDate)
    }

    // 创建表（初始化用）
    suspend fun createTableIfNotExists() {
        connectionPool.getConnection { connection ->
            val sql = """
                CREATE TABLE IF NOT EXISTS users (
                    id INT PRIMARY KEY,
                    username VARCHAR(50) NOT NULL,
                    email VARCHAR(100) NOT NULL,
                    age INT NOT NULL,
                    registration_date BIGINT NOT NULL
                )
            """.trimIndent()

            connection.createStatement().use { stmt ->
                stmt.execute(sql)
            }
        }
    }

    // 清空表
    suspend fun truncateTable() {
        connectionPool.getConnection { connection ->
            connection.createStatement().use { stmt ->
                stmt.execute("TRUNCATE TABLE users")
                connection.commit()
            }
        }
    }
}

// 插入管理器
class InsertTestManager(
    private val dbHelper: UserDatabaseHelper, private val totalUsers: Int = 100_000, private val batchSize: Int = 1000
) {
    // 1.单协程串行插入
    suspend fun testSerialInsert(startId: Int = 0): Long {
        println("开始单协程串行插入...")
        val progressRecorder = ProgressRecorder("serial_progress.txt")
        val actualStart = maxOf(startId, progressRecorder.getCurrentProgress())
        val remaining = totalUsers - actualStart

        if (remaining <= 0) {
            println("所有数据已插入完成")
            return 0
        }

        println("从ID $actualStart 开始，还需插入$remaining 条数据")

        val time = measureTimeMillis {
            for (i in actualStart until totalUsers) {
                val user = UserDataGenerator.generateUsers(i, 1).first()
                dbHelper.insertUser(user)

                // 每100条记录一次进度
                if (i % 100 == 0) {
                    progressRecorder.updateProgress(i + 1)
                    if (i % 1000 == 0) {
                        println("已插入 $i 条数据")
                    }
                }
            }
        }

        progressRecorder.updateProgress(totalUsers)
        return time
    }

    // 2.无限制并发插入
    suspend fun testUnlimitedConcurrentInsert(startId: Int = 0): Long {
        println("开始无限制并发插入...")
        val progressRecorder = ProgressRecorder("unlimited_progress.txt")
        val actualStart = maxOf(startId, progressRecorder.getCurrentProgress())
        val remaining = totalUsers - actualStart

        if (remaining <= 0) {
            println("所有数据已插入完成")
            return 0
        }

        println("从ID $actualStart 开始，还需插入 $remaining 条数据")
        val time = measureTimeMillis {
            val scope = CoroutineScope(Dispatchers.IO)
            val progress = AtomicInteger(actualStart)

            // 按批次生成任务，单不限制并发数
            val batches = (0 until (remaining + batchSize - 1) / batchSize).map { batchNum ->
                val batchStart = actualStart + batchNum * batchSize
                val batchEnd = minOf(actualStart + (batchNum + 1) * batchSize, totalUsers)

                scope.async {
                    val users = UserDataGenerator.generateUsers(batchStart, batchEnd - batchStart)
                    val inserted = dbHelper.batchInsertUsers(users)

                    // 更新进度
                    val newProgress = progress.addAndGet(inserted)
                    if (newProgress % 1000 == 0) {
                        println("已插入 $newProgress 条数据")
                    }
                    progressRecorder.updateProgress(newProgress)
                }
            }

            batches.awaitAll()
            scope.cancel()
        }

        progressRecorder.updateProgress(totalUsers)
        return time
    }

    // 3.有限制并发 + 批处理插入
    suspend fun testLimitedConcurrentInsert(
        maxConcurency: Int = 15, startId: Int = 0
    ): Long {
        println("开始有限制并发($maxConcurency) + 批处理插入...")
        val progressRecorder = ProgressRecorder("limited_progress.txt")
        val actualStart = maxOf(startId, progressRecorder.getCurrentProgress())
        val remaining = totalUsers - actualStart

        if (remaining <= 0) {
            println("所有数据已插入完成")
            return 0
        }

        println("从ID ${actualStart} 开始，还需要插入 $remaining 条数据")

        val semaphore = Semaphore(maxConcurency)
        val time = measureTimeMillis {
            val scope = CoroutineScope(Dispatchers.IO)
            val progress = AtomicInteger(actualStart)

            // 按批次生成任务，并限制并发数
            val batches = (0 until (remaining + batchSize - 1) / batchSize).map { batchNum ->
                val batchStart = actualStart + batchNum * batchSize
                val batchEnd = minOf(actualStart + (batchNum + 1) * batchSize, totalUsers)

                scope.async {
                    semaphore.withPermit {
                        val users = UserDataGenerator.generateUsers(batchStart, batchEnd - batchStart)
                        val inserted = dbHelper.batchInsertUsers(users)

                        // 更新进度
                        val newProgress = progress.addAndGet(inserted)
                        if (newProgress % 1000 == 0) {
                            println("已插入 $newProgress 条数据")
                        }
                        progressRecorder.updateProgress(newProgress)
                    }
                }
            }

            batches.awaitAll()
            scope.cancel()
        }

        progressRecorder.updateProgress(totalUsers)
        return time
    }

    // 重置所有进度记录
    fun resetAllProgress() {
        ProgressRecorder("serial_progress.txt").reset()
        ProgressRecorder("unlimited_progress.txt").reset()
        ProgressRecorder("limited_progress.txt").reset()
    }
}

fun main() = runBlocking {
    val dbUrl = "jdbc:mysql://localhost:3306/test?useSSL=false&serverTimezone=UTC"
    val dbUser = "root"
    val dbPassword = "123456"
    val maxConnections = 20
    val totalUsers = 100_000
    val batchSize = 1000

    // 初始化连接池和数据库工具
    val connectionPool = DatabaseConnectionPool(dbUrl, dbUser, dbPassword, maxConnections)
    val dbHelper = UserDatabaseHelper(connectionPool)

    try {
        // 确保表存在
        dbHelper.createTableIfNotExists()
        // 测试前清空表
        dbHelper.truncateTable()

        val testManager = InsertTestManager(dbHelper, totalUsers, batchSize)
        // 重置进度
        testManager.resetAllProgress()

        // 分别测试三种插入方式
        val serialTime = testManager.testSerialInsert()
        // 重新清空表，确保公平测试
        dbHelper.truncateTable()
        testManager.resetAllProgress()

        val unlimitedTime = testManager.testUnlimitedConcurrentInsert()
        dbHelper.truncateTable()
        testManager.resetAllProgress()

        val limitedTime = testManager.testLimitedConcurrentInsert()
        dbHelper.truncateTable()
        testManager.resetAllProgress()

        // 输出对比结果
        println("\n===== 插入性能对比 =====")
        println("总插入条数: $totalUsers")
        println("批处理大小: $batchSize")
        println("单协程串行插入耗时: ${serialTime}ms (${serialTime / 1000.0}秒)")
        println("无限制并发插入耗时: ${unlimitedTime}ms (${unlimitedTime / 1000.0}秒)")
        println("有限制并发(15)插入耗时: ${limitedTime}ms (${limitedTime / 1000.0}秒)")
        println("========================")
        println("有限制并发 vs 单协程: 快 ${(serialTime.toDouble() / limitedTime * 100).toInt()}%")
        println(
            "有限制并发 vs 无限制并发: ${if (limitedTime < unlimitedTime) "快" else "慢"} ${
                kotlin.math.abs(
                    limitedTime - unlimitedTime
                )
            }ms"
        )
    } finally {
        connectionPool.close()
    }
}